//
// Created by whitby on 8/19/22.
//


#include <iostream>
#include <string>
#include <thread>

#include "System.h"
#include "surfel_map.h"

using namespace std;

void LoadImages(const string &strPathToSequence, vector<string> &strImageLeft,
                vector<string> &strImageRight, vector<double> &timeStamps) {
    ifstream fTimes;
    string strPathTimeFile = strPathToSequence + "/times.txt";
    fTimes.open(strPathTimeFile.c_str());
    while (!fTimes.eof()) {
        string s;
        getline(fTimes, s);
        if (!s.empty()) {
            stringstream ss;
            ss << s;
            double t;
            ss >> t;
            timeStamps.push_back(t);
        }
    }

    string strPrefixLeft = strPathToSequence + "/image_0/";
    string strPrefixRight = strPathToSequence + "/image_1/";

    const size_t nTimes = timeStamps.size();
    strImageLeft.resize(nTimes);
    strImageRight.resize(nTimes);

    for (int i = 0; i < nTimes; ++i) {
        stringstream ss;
        ss << setfill('0') << setw(6) << i;
        strImageLeft[i] = strPrefixLeft + ss.str() + ".png";
        strImageRight[i] = strPrefixRight + ss.str() + ".png";
    }
}

int main(int argc, char **argv) {
    if (argc != 4) {
        cerr << endl << "Usage: ./stereo_kitti_path path_to_vocabulary path_to_settings path_to_sequence" << endl;
        return 0;
    }

    // Retrieve paths to images
    vector<string> strImageLeft;
    vector<string> strImageRight;
    vector<double> timeStamps;
    LoadImages(string(argv[3]), strImageLeft, strImageRight, timeStamps);

    const size_t nImages = strImageLeft.size();

    // Create SLAM system. It initializes all system threads and gets ready to process frames.
    ORB_SLAM2::System SLAM(argv[1], argv[2], ORB_SLAM2::System::STEREO, false);

    cout << endl << "----------" << endl;
    cout << "Start processing sequence ..." << endl;
    cout << "Images in the sequence: " << nImages << endl << endl;

    // initial ofstream
    string pose_file = "/home/whitby/CLionProjects/ORB2DenseMapping/output/pose.txt";
    ofstream pose_f;
    pose_f.open(pose_file.c_str());
    pose_f << fixed;

    string path_file = "/home/whitby/CLionProjects/ORB2DenseMapping/output/path.txt";
    ofstream path_f;
    path_f.open(path_file.c_str());
    path_f << fixed;

    string loop_file = "/home/whitby/CLionProjects/ORB2DenseMapping/output/loop.txt";
    ofstream loop_f;
    loop_f.open(loop_file.c_str());
    loop_f << fixed;

    // Main loop
    cv::Mat imLeft, imRight, imDepth;
    for (int ni = 0; ni < nImages; ++ni) {
        imLeft = cv::imread(strImageLeft[ni], cv::IMREAD_UNCHANGED);
        imRight = cv::imread(strImageRight[ni], cv::IMREAD_UNCHANGED);
        double time_stamp = timeStamps[ni];

        cout << "current time is " << time_stamp << endl;

        if (imLeft.empty()) {
            cerr << endl << "Failed to load image at: "
                 << string(strImageLeft[ni]) << endl;
            return 0;
        }

        chrono::steady_clock::time_point t1 = chrono::steady_clock::now();

        cv::Mat track_result = SLAM.TrackStereo(imLeft, imRight, time_stamp);

        vector<pair<cv::Mat, double>> result_vector;
        SLAM.GetAllPoses(result_vector);

        vector<pair<Eigen::Matrix4d, double>> allKeyFrames;
        Eigen::Matrix4d trans_form = Eigen::Matrix4d::Identity();
        for (auto &result_pair : result_vector) {
            Eigen::Matrix4d T_cw, T_wc;
            cv::Mat pose = result_pair.first;
            for (int row_i = 0; row_i < 4; ++row_i)
                for (int col_i = 0; col_i < 4; ++col_i)
                    T_cw(row_i, col_i) = pose.at<float>(row_i, col_i);

            T_wc = trans_form * T_cw.inverse();
            allKeyFrames.emplace_back(T_wc, result_pair.second);
        }

        double reference_stamp = SLAM.GetRelativePose();
        int reference_index = 0;
        double time_diff = 1e9;
        for (int i = 0; i < result_vector.size(); ++i) {
            double this_time_diff = fabs(result_vector[i].second - reference_stamp);
            if (this_time_diff < time_diff) {
                reference_index = i;
                time_diff = this_time_diff;
            }
        }
        if (time_diff < 0.01)
            printf("the reference keyframe is %d, keyframe number %d.\n", reference_index, (int)result_vector.size());

        bool is_new_keyframe = SLAM.GetKeyframeDecision();
        if (is_new_keyframe) {
            cout << "current frame is a keyframe" << endl;
        }

        Eigen::Matrix4d current_pose, publish_pose;
        for (int row_i = 0; row_i < 4; ++row_i)
            for (int col_i = 0; col_i < 4; ++col_i)
                current_pose(row_i, col_i) = track_result.at<float>(row_i, col_i);
        publish_pose = trans_form * current_pose.inverse();

        vector<pair<double, double>> loop_result;
        SLAM.GetLoopInfo(loop_result);
        vector<pair<int, int>> loop_path;

        for (int i = 0; i < loop_result.size() && i < 35; i++) {
            int first_index = -1;
            int second_index = -1;
            for (int j = 0; j < result_vector.size(); ++j) {
                if (result_vector[j].second == loop_result[i].first)
                    first_index = j;
                if (result_vector[j].second == loop_result[i].second)
                    second_index = j;
            }
            if (first_index > 0 && second_index > 0 && abs(first_index - second_index) > 10) {
                printf("the loop info %d <---> %d\n", first_index, second_index);
                loop_path.emplace_back(first_index, second_index);
            }
        }

        // Writing files
        pose_f << setprecision(9) << publish_pose(0, 0) << " " << publish_pose(0, 1) << " " << publish_pose(0, 2) << " " << publish_pose(0, 3) << " "
               << publish_pose(1, 0) << " " << publish_pose(1, 1) << " " << publish_pose(1, 2) << " " << publish_pose(1, 3) << " "
               << publish_pose(2, 0) << " " << publish_pose(2, 1) << " " << publish_pose(2, 2) << " " << publish_pose(2, 3) << endl;
        pose_f << reference_index << " " << is_new_keyframe << endl;

        path_f << allKeyFrames.size() << " ";
        for (auto &keyframe : allKeyFrames) {
            Eigen::Matrix4d keyframe_pose = keyframe.first;
            path_f << setprecision(9) << keyframe_pose(0, 0) << " " << keyframe_pose(0, 1) << " " << keyframe_pose(0, 2) << " "<< keyframe_pose(0, 3) << " "
                   << keyframe_pose(1, 0) << " " << keyframe_pose(1, 1) << " " << keyframe_pose(1, 2) << " " << keyframe_pose(1, 3) << " "
                   << keyframe_pose(2, 0) << " " << keyframe_pose(2, 1) << " " << keyframe_pose(2, 2) << " " << keyframe_pose(2, 3) << " ";
        }
        path_f << endl;

        loop_f << loop_path.size() << " ";
        for (auto &loop : loop_path) {
            loop_f << loop.first << " " << loop.second << " ";
        }
        loop_f << endl;

        printf("%d iteration end\n\n", ni);
    }

    pose_f.close();
    path_f.close();
    loop_f.close();
    cout << endl << "ORB-SLAM2 information saved!" << endl;

    return 1;
}